`timescale 1ps/1ps
module keyexp(
    input           clk
,   input           rst_n

,   input           en
,   output   reg    done

,   input  [3:0]    addr       // 扩展密钥寄存器堆
,   input  [127:0]  key        // 输入的初始密钥
,   output [127:0]  exp_key    // 扩展的密钥输出
 
    /* sbox share port */
,   output [31:0]   sbox_addr
,   input  [31:0]   sbox_out
);  

reg [3:0] state,next_state;
always @(posedge clk or negedge rst_n)begin
	if(!rst_n)
		state<= 4'd0;
	else
		state<=next_state;									
end
	
always @ (state or en) 
		case(state)	
			  4'd0:	
                if(en == 1) 
                    next_state = 4'd1;
                else 
                    next_state = 4'd0;
			  4'd1: next_state = 4'd2;
			  4'd2:	next_state = 4'd3;		
			  4'd3: next_state = 4'd4;
			  4'd4: next_state = 4'd5;
			  4'd5:	next_state = 4'd6;						             
			  4'd6: next_state = 4'd7;
			  4'd7:	next_state = 4'd8;
			  4'd8: next_state = 4'd9;
			  4'd9:	next_state = 4'd10;
			  4'd10:next_state = 4'd11;	
			  4'd11:next_state = 4'd0;
			  default: next_state = 4'd0;	  	
		endcase

reg keysel;
always @ (state) 
		case(state)	
			  4'd0:	keysel=0;
			  4'd1:	keysel=0;
			  4'd2:	keysel=1;
			  4'd3:	keysel=1;
			  4'd4:	keysel=1;
			  4'd5:	keysel=1;
			  4'd6:	keysel=1;
			  4'd7:	keysel=1;
			  4'd8:	keysel=1;
			  4'd9:	keysel=1;
			  4'd10:keysel=1;
			  4'd11:keysel=1;			  			  
			  default: keysel=0;  
		endcase

reg rndkren;
always @ (state) 
		case(state)	
			  4'd0:	rndkren=0;
			  4'd1:	rndkren=1;
			  4'd2:	rndkren=1;
			  4'd3:	rndkren=1;
			  4'd4:	rndkren=1;
			  4'd5:	rndkren=1;
			  4'd6:	rndkren=1;
			  4'd7:	rndkren=1;
			  4'd8:	rndkren=1;
			  4'd9:	rndkren=1;
			  4'd10:rndkren=1;
			  4'd11:rndkren=1;			  			  
			  default: rndkren=0;  
		endcase		

reg [3:0]wrkrfaddr;
always @ (state) 
		case(state)	
			  4'd0:	wrkrfaddr=4'd0;
			  4'd1:	wrkrfaddr=4'd0;
			  4'd2:	wrkrfaddr=4'd1;
			  4'd3:	wrkrfaddr=4'd2;
			  4'd4:	wrkrfaddr=4'd3;
			  4'd5:	wrkrfaddr=4'd4;
			  4'd6:	wrkrfaddr=4'd5;
			  4'd7:	wrkrfaddr=4'd6;
			  4'd8:	wrkrfaddr=4'd7;
			  4'd9:	wrkrfaddr=4'd8;
			  4'd10:wrkrfaddr=4'd9;
			  4'd11:wrkrfaddr=4'd10;			  			  
			  default: wrkrfaddr=4'd0;  
		endcase

always @(posedge clk or negedge rst_n) begin
	if(!rst_n)
		done <= 1'b0;
	else if (state == 4'd11)
		done <= 1'b1;
	else
		done <= 1'b0;
end

wire  rcon_en = keysel; 
wire  keyex_state = state != 4'd0; 
wire [3:0] rf_addr = keyex_state ? wrkrfaddr : addr;

wire [7:0] ron_count;
wire [31:0]xor1,xor2,xor3,xor4,xor5;
wire [31:0]subword,rotword;
wire [127:0] round_key_in,round_key_out;

// MUX
assign round_key_in = (key & {128{!keysel}}) | ({xor2,xor3,xor4,xor5} & {128{keysel}});

// REG
reg [127:0]key_reg;
always @(posedge clk or negedge rst_n) begin
    if(!rst_n)
        key_reg <= 'd0;
    else if(rndkren)
        key_reg <= round_key_in;
end
assign round_key_out = key_reg;

// ROT
assign rotword = {round_key_out[23:0],round_key_out[31:24]};

// SUB
assign sbox_addr = rotword;
assign subword   = sbox_out;

// RCON
rcon rcon0(clk,rst_n,en,rcon_en,ron_count);

// XOR
assign xor1 = {ron_count,24'h0000_00} ^ subword;
assign xor2 = round_key_out[127:96]   ^ xor1   ;
assign xor3 = round_key_out[ 95:64]   ^ xor2   ;
assign xor4 = round_key_out[ 63:32]   ^ xor3   ;
assign xor5 = round_key_out[ 31: 0]   ^ xor4   ;

// HEAP
regfile rf(clk,rst_n,rndkren,rf_addr,round_key_in,exp_key);

endmodule